﻿import numpy as np
from population import Population, Individual
from crossover_and_mutation import CrossoverAndMutation
from fitness_evaluate import FitnessEvaluate
from selection_operator import Selection


class EvolveCNN(object):
    def __init__(self, params):
        self.params = params
        self.pops = None

    def environment_selection(self, pops_parent, offs_newborn):
        v_list = []
        indi_list = []
        for indi in pops_parent:
            indi_list.append(indi)
            v_list.append(indi['acc'])
        for indi in offs_newborn:
            indi_list.append(indi)
            v_list.append(indi['acc'])

        # find the largest one's index
        max_index = np.argmax(v_list)
        selection = Selection()
        selected_index_list = selection.RouletteSelection(v_list, k=self.params['pop_size'])
        if max_index not in selected_index_list:
            first_selectd_v_list = [v_list[i] for i in selected_index_list]
            min_idx = np.argmin(first_selectd_v_list)
            selected_index_list[min_idx] = max_index

        new_pops = [indi_list[i] for i in selected_index_list]

        return new_pops

    def Select_Best_individual(self, pops):
        Best_acc = 0.0
        for individual in pops:
            if float(individual['acc']) > Best_acc:
                Best_individual = individual
                Best_acc = float(individual['acc'])

        return Best_individual

    def init_file(self):
        f_1 = open("All_train_individuals_information.txt", "w")
        f_1.seek(0)
        f_1.truncate()
        f_1.write('Evolving will be started and ALL traning individuals will be recorded' + '\n')
        f_1.write('\n')
        f_1.close()

    def individual_information_record(self, file_name, need_record_community, current_generation_, evaluated_flag=0, mutation_flag=0):
        with open(file_name, "a") as f_1:
            count_ = 0
            acc_cache = np.array([], dtype=np.int64)
            for individual_ in need_record_community:
                count_ += 1
                if evaluated_flag == 0:
                    if mutation_flag == 0:
                        f_1.write("The %sst individual information of %sst generation population is as follows:" % (count_, current_generation_) + '\n')
                    elif mutation_flag == 1:
                        f_1.write("The %sst individual information of %sst generation population after variation is as follows:" % (count_, current_generation_) + '\n')
                    for key, value in individual_.items():
                        if key == 'acc' or key == 'encode_layers':
                            continue
                        f_1.write(str(key) + ':')
                        f_1.write(str(value))
                        f_1.write('\n')
                elif evaluated_flag == 1:
                    if mutation_flag == 0:
                        f_1.write("The %sst individual information of %sst generation population after evaluation is as follows:" % (count_, current_generation_) + '\n')
                    elif mutation_flag == 1:
                        f_1.write("The %sst individual information of %sst generation population after variation and evaluation is as follows:" % (count_, current_generation_) + '\n')
                    for key, value in individual_.items():
                        if key == 'acc':
                            value = float(value)
                            acc_cache = np.append(acc_cache, value)
                        if key == 'encode_layers':
                            continue
                        f_1.write(str(key) + ':')
                        f_1.write(str(value))
                        f_1.write('\n')
                f_1.write('\n')
                f_1.flush()
            f_1.write('acc_cache' + '\n')
            f_1.write(str(acc_cache))
            f_1.write('\n')

    def do_work(self, max_gen):
        with open("All_train_individuals_information.txt", "a") as f_:
            gen_now = 1
            self.init_file()
            fitness_evaluate = FitnessEvaluate(self.params)
            population = Population(self.params)
            init_pops = population.initialize_individuals()
            f_.write(" To initialize population, information is as follows: "+ '\n')
            f_.flush()
            self.individual_information_record("All_train_individuals_information.txt", init_pops, gen_now,
                                               evaluated_flag=0, mutation_flag=0)
            # -------------------------------------------------------------------------------------------------
            print('----------------------------init_pops--------------------')
            for ones_index, ones_value in enumerate(init_pops):
                print("The %s individual of initial population is as follows:" % (ones_index + 1))
                for key, value in ones_value.items():
                    if key == 'encode_layers' or key == 'acc':
                        continue
                    print("%s:%s" % (key, value))
            print('----------------------------init_pops--------------------')
            print("population was initialized. Waiting for evaluation")
            # --------------------------------------------------------------------------------------------------
            f_.write("------------------------------------------------------------------------------------------" + '\n')
            f_.write("population was initialized. Waiting for evaluation" + '\n')
            f_.flush()
            init_pops = fitness_evaluate.train_and_eval(init_pops, gen_now, mutation_flag=0)
            # --------------------------------------------------------------------------------------------------
            print('----------------------------init_pops--------------------')
            for ones_index, ones_value in enumerate(init_pops):
                print("The %s individual of initial population after evaluation is as follows:" % (ones_index + 1))
                for key, value in ones_value.items():
                    if key == 'encode_layers':
                        continue
                    elif key == 'acc':
                        value = float(value)
                    print("%s:%s" % (key, value))
            print('----------------------------init_pops--------------------')
            # ---------------------------------------------------------------------------------------------------
            self.individual_information_record("All_train_individuals_information.txt", init_pops, gen_now,
                                               evaluated_flag=1, mutation_flag=0)
            self.pops = init_pops
            for curr_gen in range(gen_now, max_gen + 1):
                # crossover_and_mutation initial population generate mutation offsprings
                crossover_and_mutation = CrossoverAndMutation(self.pops, self.params)
                print("\n" + "The " + str(gen_now) + "st generation of cross and mutation will begin.")
                f_.write("The " + str(gen_now) + "st generation of cross and mutation will begins." + '\n')
                f_.flush()
                through_mutatation_offs = crossover_and_mutation.process()
                self.individual_information_record("All_train_individuals_information.txt", through_mutatation_offs, gen_now,
                                                   evaluated_flag=0, mutation_flag=1)
                # ---------------------------------------------------------------------------------------------------
                print('----------------------------through_mutatation_offs-------------------------')
                print("The " + str(gen_now) + "st generation offsprings information after mutation")
                for ones_index, ones_value in enumerate(through_mutatation_offs):
                    print("The %s offspring after mutation shows：" % (ones_index + 1))
                    for key, value in ones_value.items():
                        if key == 'encode_layers' or key == 'acc':
                            continue
                        print("%s:%s" % (key, value))
                print('----------------------------through_mutatation_offs-------------------------')
                # ---------------------------------------------------------------------------------------------------
                # encode mutation offsprings
                through_mutatation_offs = population.create_from_offspring(through_mutatation_offs)
                # evaluate mutation offsprings
                through_mutatation_offs = fitness_evaluate.train_and_eval(through_mutatation_offs, gen_now,
                                                                          mutation_flag=1)
                # -------------------------------------------------------------------------------------------------
                print('----------------------------through_mutatation_offs-------------------------')
                print("The " + str(gen_now) + "st  generation offsprings after mutation and evaluation")
                for ones_index, ones_value in enumerate(through_mutatation_offs):
                    print("The %s offspring information after mutation and evaluation displays：" % (ones_index + 1))
                    for key, value in ones_value.items():
                        if key == 'encode_layers':
                            continue
                        elif key == 'acc':
                            value = float(value)
                        print("%s:%s" % (key, value))
                print('----------------------------through_mutatation_offs-------------------------')
                # -------------------------------------------------------------------------------------------------
                f_.write("------------------------------------------------------------------------------------------" + '\n')
                f_.write("The " + str(gen_now) + "st generarion offsprings information through mutation and evaluation" + '\n')
                f_.flush()
                self.individual_information_record("All_train_individuals_information.txt", through_mutatation_offs, gen_now,
                                                   evaluated_flag=1, mutation_flag=1)
                print("The " + str(gen_now) + "st generation makes environmental choices and produces the next generation of individuals")
                f_.write("------------------------------------------------------------------------------------------" + '\n')
                f_.write("The " + str(gen_now) + "st generation makes environmental choices and produces the next generation of individuals" + '\n')
                f_.flush()
                # select new population from last population and offsprings through mutation
                new_pops = self.environment_selection(self.pops, through_mutatation_offs)
                """Here, the population information should be updated, such as the gene no and then to the individual id"""
                new_pops = population.create_from_offspring(new_pops)
                self.individual_information_record("All_train_individuals_information.txt", new_pops, gen_now,
                                                   evaluated_flag=0, mutation_flag=0)
                self.pops = new_pops
                # -------------------------------------------------------------------------------------------------------------
                f_.write("------------------------------------------------------------------------------------------" + '\n')
                f_.flush()
                print('----------------------------------new_pops----------------------------------')
                print("The " + str(gen_now) + "st generation population information after environmental selection is as follows:")
                for ones_index, ones_value in enumerate(new_pops):
                    print("The %sst individual from population exhibits:" % (ones_index + 1))
                    for key, value in ones_value.items():
                        if key == 'encode_layers':
                            continue
                        elif key == 'acc':
                            value = float(value)
                        print("%s:%s" % (key, value))
                print('-----------------------------------------------------------------------------')
                # -------------------------------------------------------------------------------------------------------------
                gen_now += 1
            f_.write("The individual with Best performance begin to test and train " + '\n')
            f_.flush()
            Best_individual = self.Select_Best_individual(self.pops)
            fitness_evaluate.test_end(Best_individual)

if __name__ == '__main__':
    params = {
        'prob_crossover': 0.9, #0.9
        'prob_mutation': 0.2,  # 0.2 #to test mutation component

        'BasicBlock_min_limit': 1, #3 10
        'BasicBlock_max_limit': 1,  # 即最大生成10个ResBlock
        # 'BU_outchannels_list': [64, 128, 256],
        'BU_MinOutChannels_list': 64,
        'BU_MaxOutChannels_list': 256,
        'BU_conv_type_list':[1, 2], #1 and 2 is symbol of Dilconv and Sepconv respectively
        'BU_kernel_size_list':[3, 5],

        'BU_minlimit': 1,
        'BU_maxlimit': 40,

        'PU_minlimit': 1,
        'PU_maxlimit': 4,  # 这里估计要改，调参的话，现在池化层是   size/2

        'AU_minlimit': 1,
        'AU_maxlimit': 4,
        'AU_conv_type_list': [0, 1],

        'init_max_len': 50,
        'off_max_len': 50,  #应该设大点比如15或者20
        'init_inchannels': 200,
        # 'Max_out_channels': 64,  # 用在首序列初始化和RB的内部变异

        'init_size': 32, #64  32
        'pop_size': 20,
        'max_gen': 10,
        'EPOCH': 120,  #150
        'BATCH_SIZE': 32,
        'num_class': 16,
    }
    evoCNN = EvolveCNN(params)
    evoCNN.do_work(params['max_gen'])